import asyncio
import functools
import time
import random
from typing import Callable, Dict, Any, Optional, TypeVar, cast, List
from datetime import datetime, timedelta
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.jobstores.memory import MemoryJobStore
from apscheduler.executors.asyncio import AsyncIOExecutor
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger
from apscheduler.triggers.date import DateTrigger
from loguru import logger

from utils.constants import REDIS_RETRY_ATTEMPTS, REDIS_RETRY_WAIT_SECONDS

T = TypeVar('T')


class RetryConfig:
    """
    重试配置类
    """
    def __init__(self, 
                 max_retries: int = 3, 
                 base_delay: float = 1.0, 
                 max_delay: float = 60.0, 
                 backoff_factor: float = 2.0, 
                 jitter: bool = True,
                 retry_exceptions: tuple = (Exception,),
                 retry_on_result: Optional[Callable[[Any], bool]] = None):
        """
        初始化重试配置
        :param max_retries: 最大重试次数
        :param base_delay: 基础延迟时间（秒）
        :param max_delay: 最大延迟时间（秒）
        :param backoff_factor: 退避因子
        :param jitter: 是否添加抖动
        :param retry_exceptions: 需要重试的异常类型
        :param retry_on_result: 根据结果决定是否重试的函数
        """
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.max_delay = max_delay
        self.backoff_factor = backoff_factor
        self.jitter = jitter
        self.retry_exceptions = retry_exceptions
        self.retry_on_result = retry_on_result


class RetryException(Exception):
    """
    重试失败异常
    """
    pass


def retry(*args, **kwargs):
    """
    重试装饰器（同步函数）
    
    可以作为独立装饰器使用，也可以传入RetryConfig对象
    
    示例:
        @retry(max_retries=3, base_delay=1.0)
        def my_function():
            pass
    """
    # 支持两种调用方式: @retry 或 @retry(max_retries=3)
    if len(args) == 1 and callable(args[0]):
        # 直接作为装饰器使用: @retry
        func = args[0]
        config = RetryConfig()
        
        @functools.wraps(func)
        def wrapper(*func_args, **func_kwargs):
            return _retry_sync(func, config, *func_args, **func_kwargs)
        
        return wrapper
    else:
        # 传入参数使用: @retry(max_retries=3)
        config = RetryConfig(*args, **kwargs)
        
        def decorator(func):
            @functools.wraps(func)
            def wrapper(*func_args, **func_kwargs):
                return _retry_sync(func, config, *func_args, **func_kwargs)
            return wrapper
        
        return decorator


def async_retry(*args, **kwargs):
    """
    异步重试装饰器（异步函数）
    
    可以作为独立装饰器使用，也可以传入RetryConfig对象
    
    示例:
        @async_retry(max_retries=3, base_delay=1.0)
        async def my_function():
            pass
    """
    # 支持两种调用方式: @async_retry 或 @async_retry(max_retries=3)
    if len(args) == 1 and callable(args[0]):
        # 直接作为装饰器使用: @async_retry
        func = args[0]
        config = RetryConfig()
        
        @functools.wraps(func)
        async def wrapper(*func_args, **func_kwargs):
            return await _retry_async(func, config, *func_args, **func_kwargs)
        
        return wrapper
    else:
        # 传入参数使用: @async_retry(max_retries=3)
        config = RetryConfig(*args, **kwargs)
        
        def decorator(func):
            @functools.wraps(func)
            async def wrapper(*func_args, **func_kwargs):
                return await _retry_async(func, config, *func_args, **func_kwargs)
            return wrapper
        
        return decorator


def _calculate_delay(attempt: int, base_delay: float, max_delay: float, 
                    backoff_factor: float, jitter: bool) -> float:
    """
    计算重试延迟时间
    :param attempt: 当前尝试次数
    :param base_delay: 基础延迟
    :param max_delay: 最大延迟
    :param backoff_factor: 退避因子
    :param jitter: 是否添加抖动
    :return: 延迟时间（秒）
    """
    # 指数退避: base_delay * (backoff_factor ^ (attempt - 1))
    delay = base_delay * (backoff_factor ** (attempt - 1))
    
    # 限制最大延迟
    delay = min(delay, max_delay)
    
    # 添加抖动（±20%）
    if jitter:
        delay *= random.uniform(0.8, 1.2)
    
    return delay


def _retry_sync(func: Callable[..., T], config: RetryConfig, 
                *args: Any, **kwargs: Any) -> T:
    """
    同步函数重试逻辑
    """
    last_exception = None
    
    for attempt in range(config.max_retries + 1):  # +1 是因为首次尝试不计入重试次数
        try:
            result = func(*args, **kwargs)
            
            # 检查是否需要基于结果重试
            if config.retry_on_result and config.retry_on_result(result):
                logger.warning(f"函数 {func.__name__} 返回需要重试的结果，将进行重试，尝试次数: {attempt + 1}/{config.max_retries}")
                if attempt < config.max_retries:
                    delay = _calculate_delay(attempt + 1, config.base_delay, 
                                           config.max_delay, config.backoff_factor, 
                                           config.jitter)
                    logger.info(f"等待 {delay:.2f} 秒后重试...")
                    time.sleep(delay)
                    continue
                else:
                    logger.error(f"函数 {func.__name__} 达到最大重试次数")
                    raise RetryException(f"函数 {func.__name__} 达到最大重试次数: {config.max_retries}")
            
            # 重试成功，记录日志
            if attempt > 0:
                logger.info(f"函数 {func.__name__} 重试成功，尝试次数: {attempt}/{config.max_retries}")
            
            return result
        
        except config.retry_exceptions as e:
            last_exception = e
            
            # 如果不是最后一次尝试，继续重试
            if attempt < config.max_retries:
                delay = _calculate_delay(attempt + 1, config.base_delay, 
                                       config.max_delay, config.backoff_factor, 
                                       config.jitter)
                logger.warning(f"函数 {func.__name__} 执行失败，将在 {delay:.2f} 秒后重试，错误: {str(e)}，尝试次数: {attempt + 1}/{config.max_retries}")
                time.sleep(delay)
            else:
                logger.error(f"函数 {func.__name__} 重试失败，达到最大重试次数，最后错误: {str(e)}")
    
    # 所有重试都失败，抛出异常
    if last_exception:
        raise RetryException(f"函数 {func.__name__} 重试失败: {str(last_exception)}") from last_exception
    
    # 理论上不会到达这里，但为了类型安全
    raise RetryException(f"函数 {func.__name__} 执行失败")


async def _retry_async(func: Callable[..., T], config: RetryConfig, 
                      *args: Any, **kwargs: Any) -> T:
    """
    异步函数重试逻辑
    """
    last_exception = None
    
    for attempt in range(config.max_retries + 1):  # +1 是因为首次尝试不计入重试次数
        try:
            result = await func(*args, **kwargs)
            
            # 检查是否需要基于结果重试
            if config.retry_on_result and config.retry_on_result(result):
                logger.warning(f"异步函数 {func.__name__} 返回需要重试的结果，将进行重试，尝试次数: {attempt + 1}/{config.max_retries}")
                if attempt < config.max_retries:
                    delay = _calculate_delay(attempt + 1, config.base_delay, 
                                           config.max_delay, config.backoff_factor, 
                                           config.jitter)
                    logger.info(f"等待 {delay:.2f} 秒后重试...")
                    await asyncio.sleep(delay)
                    continue
                else:
                    logger.error(f"异步函数 {func.__name__} 达到最大重试次数")
                    raise RetryException(f"异步函数 {func.__name__} 达到最大重试次数: {config.max_retries}")
            
            # 重试成功，记录日志
            if attempt > 0:
                logger.info(f"异步函数 {func.__name__} 重试成功，尝试次数: {attempt}/{config.max_retries}")
            
            return result
        
        except config.retry_exceptions as e:
            last_exception = e
            
            # 如果不是最后一次尝试，继续重试
            if attempt < config.max_retries:
                delay = _calculate_delay(attempt + 1, config.base_delay, 
                                       config.max_delay, config.backoff_factor, 
                                       config.jitter)
                logger.warning(f"异步函数 {func.__name__} 执行失败，将在 {delay:.2f} 秒后重试，错误: {str(e)}，尝试次数: {attempt + 1}/{config.max_retries}")
                await asyncio.sleep(delay)
            else:
                logger.error(f"异步函数 {func.__name__} 重试失败，达到最大重试次数，最后错误: {str(e)}")
    
    # 所有重试都失败，抛出异常
    if last_exception:
        raise RetryException(f"异步函数 {func.__name__} 重试失败: {str(last_exception)}") from last_exception
    
    # 理论上不会到达这里，但为了类型安全
    raise RetryException(f"异步函数 {func.__name__} 执行失败")


class TaskScheduler:
    """
    任务调度器（基于APScheduler）
    """
    _instance = None
    _lock = asyncio.Lock()
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def __init__(self):
        if not hasattr(self, '_initialized'):
            # 配置调度器
            jobstores = {
                'default': MemoryJobStore()
            }
            executors = {
                'default': AsyncIOExecutor()
            }
            job_defaults = {
                'coalesce': True,  # 合并错过的任务
                'max_instances': 1,  # 每个任务的最大实例数
                'misfire_grace_time': 60  # 任务错过后的宽限时间（秒）
            }
            
            # 创建异步调度器
            self.scheduler = AsyncIOScheduler(
                jobstores=jobstores,
                executors=executors,
                job_defaults=job_defaults,
                timezone='Asia/Shanghai'  # 使用上海时区
            )
            
            self._initialized = False
            self._jobs = {}  # 存储任务ID和任务信息的映射
            
            # 添加任务执行监听器
            self.scheduler.add_listener(self._job_listener, mask=1|2|4|8)  # 所有事件
    
    async def initialize(self):
        """
        初始化调度器
        """
        async with self._lock:
            if not self._initialized:
                try:
                    self.scheduler.start()
                    self._initialized = True
                    logger.info("任务调度器初始化成功")
                    
                    # 添加健康检查任务
                    await self.add_interval_task(
                        func=self._health_check,
                        seconds=300,  # 5分钟一次
                        task_id="scheduler_health_check",
                        task_name="调度器健康检查"
                    )
                    
                except Exception as e:
                    logger.error(f"任务调度器初始化失败: {str(e)}")
                    raise
    
    async def shutdown(self):
        """
        关闭调度器
        """
        if self._initialized:
            self.scheduler.shutdown()
            self._initialized = False
            logger.info("任务调度器已关闭")
    
    async def add_interval_task(self, 
                               func: Callable, 
                               seconds: int = 0,
                               minutes: int = 0,
                               hours: int = 0,
                               days: int = 0,
                               task_id: str = None,
                               task_name: str = None,
                               **kwargs) -> str:
        """
        添加间隔任务
        :param func: 要执行的函数
        :param seconds: 间隔秒数
        :param minutes: 间隔分钟数
        :param hours: 间隔小时数
        :param days: 间隔天数
        :param task_id: 任务ID（可选，不提供则自动生成）
        :param task_name: 任务名称（可选）
        :return: 任务ID
        """
        if not self._initialized:
            await self.initialize()
        
        task_id = task_id or f"interval_{int(time.time())}_{random.randint(1000, 9999)}"
        task_name = task_name or f"间隔任务_{task_id}"
        
        try:
            # 移除已存在的同名任务
            self._remove_job(task_id)
            
            # 添加新任务
            job = self.scheduler.add_job(
                func=func,
                trigger=IntervalTrigger(
                    seconds=seconds,
                    minutes=minutes,
                    hours=hours,
                    days=days
                ),
                id=task_id,
                name=task_name,
                replace_existing=True,
                **kwargs
            )
            
            self._jobs[task_id] = {
                'task_id': task_id,
                'task_name': task_name,
                'type': 'interval',
                'func': func.__name__ if hasattr(func, '__name__') else str(func),
                'created_at': datetime.now(),
                'job': job
            }
            
            logger.info(f"添加间隔任务成功: {task_name} (ID: {task_id})")
            return task_id
        
        except Exception as e:
            logger.error(f"添加间隔任务失败: {str(e)}")
            raise
    
    async def add_cron_task(self, 
                           func: Callable,
                           cron_expression: str = None,
                           year: int = None,
                           month: int = None,
                           day: int = None,
                           hour: int = None,
                           minute: int = None,
                           second: int = None,
                           task_id: str = None,
                           task_name: str = None,
                           **kwargs) -> str:
        """
        添加定时任务
        :param func: 要执行的函数
        :param cron_expression: cron表达式（可选，例如: "0 * * * *" 表示每小时执行一次）
        :param year: 年份
        :param month: 月份
        :param day: 日期
        :param hour: 小时
        :param minute: 分钟
        :param second: 秒
        :param task_id: 任务ID（可选）
        :param task_name: 任务名称（可选）
        :return: 任务ID
        """
        if not self._initialized:
            await self.initialize()
        
        task_id = task_id or f"cron_{int(time.time())}_{random.randint(1000, 9999)}"
        task_name = task_name or f"定时任务_{task_id}"
        
        try:
            # 移除已存在的同名任务
            self._remove_job(task_id)
            
            if cron_expression:
                # 使用cron表达式
                trigger = CronTrigger.from_crontab(cron_expression)
            else:
                # 使用参数构建cron触发器
                trigger = CronTrigger(
                    year=year,
                    month=month,
                    day=day,
                    hour=hour,
                    minute=minute,
                    second=second
                )
            
            # 添加新任务
            job = self.scheduler.add_job(
                func=func,
                trigger=trigger,
                id=task_id,
                name=task_name,
                replace_existing=True,
                **kwargs
            )
            
            self._jobs[task_id] = {
                'task_id': task_id,
                'task_name': task_name,
                'type': 'cron',
                'func': func.__name__ if hasattr(func, '__name__') else str(func),
                'created_at': datetime.now(),
                'job': job
            }
            
            logger.info(f"添加定时任务成功: {task_name} (ID: {task_id})")
            return task_id
        
        except Exception as e:
            logger.error(f"添加定时任务失败: {str(e)}")
            raise
    
    async def add_once_task(self, 
                           func: Callable,
                           run_date: datetime = None,
                           seconds: int = None,
                           task_id: str = None,
                           task_name: str = None,
                           **kwargs) -> str:
        """
        添加一次性任务
        :param func: 要执行的函数
        :param run_date: 执行时间（可选）
        :param seconds: 延迟秒数（可选，如果不提供run_date则使用）
        :param task_id: 任务ID（可选）
        :param task_name: 任务名称（可选）
        :return: 任务ID
        """
        if not self._initialized:
            await self.initialize()
        
        task_id = task_id or f"once_{int(time.time())}_{random.randint(1000, 9999)}"
        task_name = task_name or f"一次性任务_{task_id}"
        
        try:
            # 移除已存在的同名任务
            self._remove_job(task_id)
            
            # 确定执行时间
            if run_date:
                trigger_date = run_date
            elif seconds is not None:
                trigger_date = datetime.now() + timedelta(seconds=seconds)
            else:
                trigger_date = datetime.now()
            
            # 添加新任务
            job = self.scheduler.add_job(
                func=func,
                trigger=DateTrigger(run_date=trigger_date),
                id=task_id,
                name=task_name,
                replace_existing=True,
                **kwargs
            )
            
            self._jobs[task_id] = {
                'task_id': task_id,
                'task_name': task_name,
                'type': 'once',
                'func': func.__name__ if hasattr(func, '__name__') else str(func),
                'created_at': datetime.now(),
                'run_date': trigger_date,
                'job': job
            }
            
            logger.info(f"添加一次性任务成功: {task_name} (ID: {task_id}), 执行时间: {trigger_date}")
            return task_id
        
        except Exception as e:
            logger.error(f"添加一次性任务失败: {str(e)}")
            raise
    
    def remove_task(self, task_id: str):
        """
        移除任务
        :param task_id: 任务ID
        """
        try:
            self._remove_job(task_id)
            if task_id in self._jobs:
                del self._jobs[task_id]
            logger.info(f"移除任务成功: {task_id}")
        except Exception as e:
            logger.error(f"移除任务失败: {str(e)}")
            raise
    
    def _remove_job(self, task_id: str):
        """
        内部方法：移除任务（不抛异常）
        """
        try:
            self.scheduler.remove_job(task_id)
        except Exception as e:
            # 任务不存在时忽略错误
            logger.debug(f"任务不存在或已移除: {task_id}, 错误: {str(e)}")
    
    def get_job(self, task_id: str) -> Dict[str, Any]:
        """
        获取任务信息
        :param task_id: 任务ID
        :return: 任务信息字典
        """
        job_info = self._jobs.get(task_id)
        if job_info:
            # 获取最新的任务状态
            job = self.scheduler.get_job(task_id)
            if job:
                job_info['next_run_time'] = job.next_run_time
                job_info['pending'] = job.pending
            return job_info
        return None
    
    def get_all_jobs(self) -> List[Dict[str, Any]]:
        """
        获取所有任务信息
        :return: 任务信息列表
        """
        jobs = []
        for task_id, job_info in self._jobs.items():
            job = self.scheduler.get_job(task_id)
            if job:
                job_info_copy = job_info.copy()
                job_info_copy['next_run_time'] = job.next_run_time
                job_info_copy['pending'] = job.pending
                jobs.append(job_info_copy)
        return jobs
    
    def _job_listener(self, event):
        """
        任务执行监听器
        """
        try:
            if event.code == 1:  # EVENT_JOB_ADDED
                logger.debug(f"任务已添加: {event.job_id}")
            elif event.code == 2:  # EVENT_JOB_REMOVED
                logger.debug(f"任务已移除: {event.job_id}")
            elif event.code == 4:  # EVENT_JOB_MODIFIED
                logger.debug(f"任务已修改: {event.job_id}")
            elif event.code == 8:  # EVENT_JOB_EXECUTED
                job_info = self._jobs.get(event.job_id)
                if job_info:
                    logger.debug(f"任务执行成功: {job_info['task_name']} (ID: {event.job_id}), 执行时间: {event.scheduled_run_time}")
            elif event.code == 16:  # EVENT_JOB_ERROR
                job_info = self._jobs.get(event.job_id)
                if job_info:
                    logger.error(f"任务执行失败: {job_info['task_name']} (ID: {event.job_id}), 错误: {str(event.exception)}")
            elif event.code == 32:  # EVENT_JOB_MISSED
                job_info = self._jobs.get(event.job_id)
                if job_info:
                    logger.warning(f"任务执行错过: {job_info['task_name']} (ID: {event.job_id}), 计划时间: {event.scheduled_run_time}")
        except Exception as e:
            logger.error(f"任务监听器出错: {str(e)}")
    
    async def _health_check(self):
        """
        调度器健康检查
        """
        try:
            jobs_count = len(self.get_all_jobs())
            logger.debug(f"调度器健康检查: 运行中的任务数={jobs_count}")
            
            # 检查是否有长时间未执行的任务
            now = datetime.now()
            for job_info in self.get_all_jobs():
                if job_info.get('next_run_time') and job_info['next_run_time'] < now - timedelta(minutes=5):
                    logger.warning(f"检测到可能延迟的任务: {job_info['task_name']}, 下次运行时间: {job_info['next_run_time']}")
        
        except Exception as e:
            logger.error(f"调度器健康检查失败: {str(e)}")


# 创建全局任务调度器实例
task_scheduler = TaskScheduler()


# 预设的向量数据库重试配置
VECTOR_DB_RETRY_CONFIG = RetryConfig(
    max_retries=3,
    base_delay=1.0,
    max_delay=10.0,
    backoff_factor=2.0,
    jitter=True,
    # 这里可以根据实际需要定义需要重试的异常类型
    retry_exceptions=(Exception,)
)


# 用于向量数据库操作的重试装饰器
def vector_db_retry(func=None):
    """
    专为向量数据库操作设计的重试装饰器
    """
    if func is None:
        def decorator(func):
            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                return _retry_sync(func, VECTOR_DB_RETRY_CONFIG, *args, **kwargs)
            return wrapper
        return decorator
    else:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return _retry_sync(func, VECTOR_DB_RETRY_CONFIG, *args, **kwargs)
        return wrapper


# 异步向量数据库操作重试装饰器
async def async_vector_db_retry(func):
    """
    异步向量数据库操作重试装饰器
    """
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        return await _retry_async(func, VECTOR_DB_RETRY_CONFIG, *args, **kwargs)
    return wrapper
